package leetcode.hash;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

public class UnionFindSet {
    //包装结点
    public static class Element<V> {
        public V value;

        public Element(V value) {
            this.value = value;

        }
    }

    public static class unionFindSet<V> {

        public HashMap<V, Element<V>> elementHashMap;//值到结点的映射
        public HashMap<Element<V>, Element<V>> fatherMap;//k:元素 V:该元素的父元素
        public HashMap<Element<V>, Integer> rankMap;//key:某集合代表元素 V:集合大小

        public unionFindSet(List<V> list) { //初始化，一个列表里，每个元素都单独一个集合

            elementHashMap = new HashMap<>();
            fatherMap = new HashMap<>();
            rankMap = new HashMap<>();

            for (V value : list) {
                Element<V> element = new Element<>(value);
                elementHashMap.put(value, element);
                fatherMap.put(element, element);
                rankMap.put(element, 1);
            }

        }

        private Element<V> findHead(Element<V> element) {

            Stack<Element<V>> path = new Stack<>();
            //如果传入节点的父节点不是该节点本身，将该节点压入栈，将指针指向其父节点循环判断直到找到头
            while (element != fatherMap.get(element)) {
                path.push(element);
                element = fatherMap.get(element);
            }
            //如果path不为空将栈中所有结点的父节点改为”头“
            while (!path.isEmpty()) {
                fatherMap.put(path.pop(), element);
            }

            return element;
        }

        //判断头节点是否相同
        public boolean isSameSet(V a, V b) {
            if (elementHashMap.containsKey(a) && elementHashMap.containsKey(b)) {
                return findHead(elementHashMap.get(a)) == findHead(elementHashMap.get(b));
            }
            return false;
        }

        public void union(V a, V b) {
            if (elementHashMap.containsKey(a) && elementHashMap.containsKey(b)) {
                Element<V> aF = findHead(elementHashMap.get(a));
                Element<V> bF = findHead(elementHashMap.get(a));
                if (aF != bF) {
                    Element<V> big = rankMap.get(aF) >= rankMap.get(bF) ? aF : bF;
                    Element<V> small = big == aF ? bF : aF;
                    fatherMap.put(small, big);
                    rankMap.put(big, rankMap.get(aF) + rankMap.get(bF));
                    rankMap.remove(small);
                }
            }
        }
    }

}
